import numpy as np
import matplotlib.pyplot as plt
from scipy.linalg import block_diag
import math
def demo3_minimum_snap_close_form():
    # 清理环境
    # condition
    waypts = np.array([[0, 0,0],[10,10,40]]).T
    v0 = np.array([0, 0,4])
    a0 = np.array([0, 0,0])
    v1 = np.array([0, 0,0])
    a1 = np.array([0, -4,0])
    T = 10
    ts = arrangeT(waypts, T)
    n_order = 5

    #trajectory plan
    polys_x = minimum_snap_single_axis_close_form(waypts[0, :], ts, n_order, v0[0], a0[0], v1[0], a1[0])
    polys_y = minimum_snap_single_axis_close_form(waypts[1, :], ts, n_order, v0[1], a0[1], v1[1], a1[1])
    polys_z = minimum_snap_single_axis_close_form(waypts[2, :], ts, n_order, v0[2], a0[2], v1[2], a1[2])

    # result show
    fig=plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    ax.plot3D(waypts[0, :], waypts[1, :], waypts[2,:],'*r')
    ax.plot(waypts[0, :], waypts[1, :], waypts[2,:], 'b--')
    # ax.title('Minimum Snap Trajectory')
    color = ['g', 'r', 'c']

    for i in range(polys_x.shape[1]):
        tt = np.arange(ts[i], ts[i + 1], 0.01)
        xx = polys_vals(polys_x, ts, tt, 0)
        yy = polys_vals(polys_y, ts, tt, 0)
        zz = polys_vals(polys_z, ts, tt, 0)
        ax.plot3D(xx, yy, zz)
        # plt.plot(xx, yy, color[i % 3])
        xx1 = polys_vals(polys_x, ts, tt, 1)
        yy1 = polys_vals(polys_y, ts, tt, 1)
        zz1 = polys_vals(polys_z, ts, tt, 1)
        v=np.sqrt(xx1**2+yy1**2+zz1**2)
        print("v",v)
        # plt.plot(tt, xx1, color[i % 3])
        xx2 = polys_vals(polys_x, ts, tt, 2)
        yy2 = polys_vals(polys_y, ts, tt, 2)
        zz2 = polys_vals(polys_z, ts, tt, 2)
        a = np.sqrt(xx2 ** 2 + yy2 ** 2 + zz2 ** 2)
        print("a",a)
        # plt.plot(tt, xx2, color[i % 3])
    # 调整子图间距
    plt.axis('equal')
    plt.show()


def minimum_snap_single_axis_close_form(wayp, ts, n_order, v0, a0, v1, a1):
    # 参数数量
    n_coef = n_order + 1
    # 线段数量
    n_poly = len(wayp) - 1

    q_i_cost = []
    for i in range(n_poly):
        q_i_cost.append(computeQ(n_order, 3, ts[i], ts[i + 1]))

    Q_all = block_diag(*q_i_cost)
    #compute Tk
    tk = np.zeros((n_poly + 1, n_coef))
    for i in range(n_coef):
        tk[:, i] = ts ** i

    # compute A
    n_continuous = 3  # 1:p  2:pv  3:pva
    A = np.zeros((n_continuous * 2 * n_poly, n_coef * n_poly))

    for i in range(n_poly): # 没有个线段
        for j in range(n_continuous): # 每一个求导
            for k in range(j, n_coef): # 每一个参数
                t1 = tk[i, k - j]
                t2 = tk[i + 1, k - j]

                A[n_continuous * 2 * i + j, n_coef * i + k] = math.prod(np.arange(k-j+1, k+1)) * t1
                A[n_continuous * 2 * i + n_continuous + j, n_coef * i + k] = math.prod(np.arange(k-j+1, k+1)) * t2

    # compute M
    M = np.zeros((n_poly * 2 * n_continuous, n_continuous * (n_poly + 1)))
    for i in range(1, n_poly * 2 + 1):
        j = i//2+1
        rbeg = n_continuous * (i - 1)
        cbeg = n_continuous * (j-1)
        M[rbeg:rbeg + n_continuous, cbeg:cbeg + n_continuous] = np.eye(n_continuous)

    # compute C
    num_d = n_continuous * (n_poly + 1)
    C = np.eye(num_d)
    df = np.concatenate((wayp, [v0], [a0], [v1], [a1])).reshape(-1, 1)  # fix all pos + start va + end va
    fix_idx = np.array([i for i in range(0, num_d, 3)] + [1, 2, num_d - 2, num_d - 1])
    free_idx = np.setdiff1d(np.arange(num_d), fix_idx)
    C = np.hstack((C[:, fix_idx], C[:, free_idx]))

    AiMC = np.linalg.inv(A) @ M @ C
    R = AiMC.T @ Q_all @ AiMC
    n_fix = len(fix_idx)
    Rff = R[:n_fix, :n_fix]
    Rpp = R[n_fix:, n_fix:]
    Rfp = R[:n_fix, n_fix:]
    Rpf = R[n_fix:, :n_fix]
    dp = -np.linalg.inv(Rpp) @ Rfp.T @ df
    p = AiMC @ np.vstack((df, dp))
    polys = p.reshape(n_poly, n_coef).T
    return polys



def arrangeT(waypts, T):
    # 计算相邻路径点之间的差值
    x = waypts[:, 1:] - waypts[:, :-1]
    # 计算路径点之间的欧几里得距离
    dist = np.sum(x ** 2, axis=0) ** 0.5

    # 计算总距离
    total_dist = np.sum(dist)

    # 计算时间缩放因子
    k = T / total_dist

    # 计算时间序列
    ts = np.concatenate(([0], np.cumsum(dist * k)))

    return ts



def computeQ(n, r, t1, t2):

    # 初始化 Q 矩阵
    Q = np.zeros((n + 1, n + 1))

    # 填充 Q 矩阵
    for i in range(r, n + 1):  # 行
        for j in range(i, n + 1):  # 列
            k1 = i - r
            k2 = j - r
            k = k1 + k2 + 1
            Q[i, j] = (math.factorial(i) / math.factorial(k1)) * (math.factorial(j) / math.factorial(k2)) * (
                        pow(t2, k) - pow(t1, k)) / k
            Q[j, i] = Q[i, j]
    return Q


def poly_val(polys, t, r):
    # 计算多项式的值
    n = len(polys) - 1
    val = 0
    for i in range(n + 1):
        val += polys[i] * (t ** (i + r))
    return val

def polys_vals(polys, ts, tt, r):
    idx = 0
    N = len(tt)
    vals = np.zeros(N)
    for i in range(N):
        t = tt[i]
        if t < ts[idx]:
            vals[i] = 0
        else:
            while idx < len(ts) - 1 and t > ts[idx + 1] + 0.0001:
                idx += 1
            vals[i] = poly_val(polys[:, idx], t, r)
    return vals

def poly_val(poly, t, r):
    val = 0
    n = len(poly) - 1
    if r <= 0:
        for i in range(n + 1):
            val += poly[i] * (t ** i)
    else:
        for i in range(r, n + 1):
            a = poly[i] * np.prod(np.arange(i - r + 1, i + 1)) * (t ** (i - r))
            val += a
    return val

# 调用主函数
demo3_minimum_snap_close_form()